from mylangchain.agents import Tool
from langchain_core.output_parsers import JsonOutputParser
from langchain_core.runnables import RunnablePassthrough, RunnableLambda
from langchain_ollama import ChatOllama
from langchain_core.prompts import PromptTemplate
from operator import itemgetter
import requests
import re
from mylangchain.agents import initialize_agent, AgentType

parse=JsonOutputParser

# 1. 定义一个自定义工具
def get_weather(location: str) -> str:
    """调用天气 API 获取天气信息"""
    api_url = f"https://api.weatherapi.com/v1/current.json?key=264efe8419d84d64ab371844251707&q={location.strip()}"
    response = requests.get(api_url)
    data = response.json()
    weather = data["current"]["condition"]["text"]
    return f"{location} 天气API JSON返回值为：{data}"


weather_tool = Tool(
    name="WeatherFetcher",
    func=get_weather,
    description="用于获取指定城市天气信息"
)

llm = ChatOllama(base_url="http://10.0.2.114:11434", model="deepseek-r1:7b", temperature=1.0)

prompt = PromptTemplate(input_variables=["context", "question"], template="""
根据提供的上下文内容，回答用户的问题：

要求如下：
1.如果上下文跟问题有关系，必须参考上下文内容
2.如果上下文跟问题没有关系，则忽略上下文内容
3.详细并且周全的回答问题
4.如果上下文中包含一些英文，必须翻译为中文

上下文内容：
{context}

问题：
{question}
"""
)

extract_prompt = PromptTemplate.from_template("""
从用户输入中提取城市名，只返回城市名，不要解释。
用户输入：{question}
""")

extract_chain = extract_prompt | llm | RunnableLambda(lambda x: clean_reasoning(x.content).strip())


def clean_reasoning(text: str) -> str:
    """去掉 <think>...</think> 及其内容"""
    return re.sub(r"<think>.*?</think>", "", text, flags=re.DOTALL).strip()

chain_weather = (
        {"context": itemgetter("question") | extract_chain | RunnableLambda(lambda city: weather_tool.run(city)),
         "question": itemgetter("question")}
        | prompt
        | llm
)

for chunk in chain_weather.stream({"question": "沈阳的天气怎么样"}):
    print(chunk.content, end="", flush=True)
